import hashlib
import os
import ctypes
import struct
import shutil
import traceback

from my_api import *


class Encrypt():
    def __init__(self, path, logger=None):
        self.output_path = Path(path, "none")
        self.input_path = Path(path, "none")
        self.ENC_BLOCK_LEN = 512
        self.logger = logger

    def combinint(self, data, sizi):
        s = {}
        num = 0
        for i in range(0, sizi, 4):
            m = 0x00000000
            for p in range(0, 4):
                q = (data[i+p] << ((p)*8))
                m = q | m
            s[num] = m
            num = num + 1
        return s

    def efuse_id_handle(self, data):
        uid = ctypes.c_uint(data)
        uid.value = uid.value << 2
        uid.value = ~uid.value
        uid.value = uid.value + 1
        return uid.value

    def number_to_str(self, data):
        enstr = b""
        for dataitm in data:
            for i in range(0, 4):
                enstr = enstr + ((data[dataitm] >> (8 * i)) & 0x000000ff).to_bytes(length=1, byteorder='little')
        return enstr

    def encrypt_secret_key(self, secret_key, efuse_id = False):
        kc = [0x23456600, 0x4acc6600, 0x8a000000, 0x9b123689]
        magic = 0x55447ab9
        i = 0
        if len(secret_key) == 16:
            cust_keynum = self.combinint(secret_key, 16)
            comdict = {}
            comdict[0] = magic
            comdict[1] = cust_keynum[0]
            comdict[2] = cust_keynum[1]
            comdict[3] = cust_keynum[2]
            comdict[4] = cust_keynum[3]
            if not efuse_id:
                enscrt = self.enbteas(comdict, 5, kc)
                return struct.pack("L" * 5, *[enscrt[k] for k in enscrt]) + struct.pack("L", 0xbabeface)
            else:
                comdict[5] = self.efuse_id_handle(efuse_id)
                enscrt = self.enbteas(comdict, 6, kc)
                enscrtStr = self.number_to_str(enscrt)
                return enscrtStr
        else:
            return False

    def handle_left_shift(self, data, num):
        return (data << num) & 0x00000000ffffffff

    def enbteas(self, v, n, key):
        DELTA = ctypes.c_uint(0x9e3779b9)
        y = ctypes.c_uint(0)
        z = ctypes.c_uint(0)
        sum = ctypes.c_uint(0)
        p = ctypes.c_uint(0)
        rounds = 0
        e = ctypes.c_uint(0)

        if n > 0:
            rounds = 6 + (52 // n)
            sum.value = 0
            z.value = v[n - 1]
            while rounds > 0:
                sum.value = sum.value + DELTA.value
                e.value = (sum.value >> 2) & 3
                p = 0
                for p in range(0, n - 1):
                    y.value = v[p + 1]
                    z.value = v[p] + (((z.value >> 5 ^ self.handle_left_shift(y.value, 2)) + (y.value >> 3 ^ self.handle_left_shift(z.value, 4))) ^ ((sum.value ^ y.value) + (key[(p & 3) ^ e.value] ^ z.value)))
                    v[p] = z.value
                y.value = v[0]
                p = p + 1
                z.value = v[n - 1] + (((z.value >> 5 ^ self.handle_left_shift(y.value, 2)) + (y.value >> 3 ^ self.handle_left_shift(z.value, 4))) ^ ((sum.value ^ y.value) + (key[(p & 3) ^ e.value] ^ z.value)))
                v[n - 1] = z.value
                # print("\r\nbtra num :" + str(rounds))
                # p = 0
                # for p in range(0, n):
                #     print(str(v[p]) + "")
                rounds = rounds - 1
        return v

    def enbtea(self, v, n, key):
        DELTA = ctypes.c_uint(0x9e3779b9)
        y = ctypes.c_uint(0)
        z = ctypes.c_uint(0)
        sum = ctypes.c_uint(0)
        p = ctypes.c_uint(0)
        rounds = 0
        e = ctypes.c_uint(0)

        if n > 0:
            rounds = 6 + 52 // n
            sum.value = 0
            z.value = v[n - 1]
            while rounds > 0:
                sum.value = sum.value + DELTA.value
                e.value = (sum.value >> 2) & 3
                p = 0
                for p in range(0, n - 1):
                    y.value = v[p + 1]
                    z.value = v[p] + (((z.value >> 5 ^ self.handle_left_shift(y.value, 2)) + (y.value >> 3 ^ self.handle_left_shift(z.value, 4))) ^ ((sum.value ^ y.value) + (key[(p & 3) ^ e.value].value ^ z.value)))
                    v[p] = z.value
                y.value = v[0]
                p = p + 1
                z.value = v[n - 1] + (((z.value >> 5 ^ self.handle_left_shift(y.value, 2)) + (y.value >> 3 ^ self.handle_left_shift(z.value, 4))) ^ ((sum.value ^ y.value) + (key[(p & 3) ^ e.value].value ^ z.value)))
                v[n - 1] = z.value
                p = 0
                rounds = rounds - 1
        return v

    def form_key_get_kc(self, key):
        kc = []
        kc.append(ctypes.c_uint(0))
        kc.append(ctypes.c_uint(0))
        kc.append(ctypes.c_uint(0))
        kc.append(ctypes.c_uint(0))
        j = 0
        m1 = hashlib.md5()
        m1.update(key)
        bb_o = m1.hexdigest()
        bb = []
        for i in range(0, len(bb_o), 2):
            bb.append(int(bb_o[i:i+2], 16))

        if len(bb) != 16:
            return False
        for i in range(0, 16):
            d = ctypes.c_uint(bb[i])
            if (i % 4) == 0:
                j = 0
            kc[i//4] = ctypes.c_uint(kc[i//4].value | (d.value << j))
            j += 8
        return kc

    def encrypt_file(self, kc, file_path):
        excess_str = b""
        file_name = os.path.basename(file_path)
        enbtea_data = []
        with open(self.input_path.with_name(file_path), "rb") as f:
            all_the_text = f.read()

        enc_len = len(all_the_text)
        dummy_len = enc_len % 512
        if dummy_len > 0:
            if dummy_len == 511:
                all_the_text = all_the_text + b"\x0a"
            else:
                all_the_text = all_the_text + b"\x20" * (511 - dummy_len) + b"\x0a"
            enc_len += (512 - dummy_len)
        uint_num = enc_len//4
        if enc_len % 4 != 0:
            print(enc_len)

        data_p = struct.unpack("I"*uint_num, all_the_text)
        i = 0
        flag_num = 0
        for i in range(0, enc_len//self.ENC_BLOCK_LEN):
            temp_data = []
            for j in range(i*(self.ENC_BLOCK_LEN // 4), i*(self.ENC_BLOCK_LEN // 4)+self.ENC_BLOCK_LEN // 4):
                temp_data.append(data_p[j])
            enbtea_data.append(self.enbtea(temp_data, self.ENC_BLOCK_LEN//4, kc))
            flag_num = flag_num + 1
        if flag_num > 0:
            i = i+1
        if (enc_len % self.ENC_BLOCK_LEN) > 3:
            temp_data = []
            for j in range(i*(self.ENC_BLOCK_LEN // 4), i*(self.ENC_BLOCK_LEN // 4) + (enc_len % self.ENC_BLOCK_LEN)//4):
                temp_data.append(data_p[j])
            enbtea_data.append(self.enbtea(temp_data, (enc_len % self.ENC_BLOCK_LEN)//4, kc))

        enc_data = bytearray()
        for i in enbtea_data:
            for j in i:
                enc_data.extend(struct.pack("I", j))
        # enc_data.extend(excess_str)
        with open(self.output_path.with_name(file_name+"e"), "wb") as f:
            f.write(enc_data)

    def encrypt_file_list(self, key, dir, file_list):
        if type(key) is bytes:
            pass
        else:
            self.logger.errorr('key is not bytes')
            return False
        try:
            if os.path.exists(self.output_path.parent):
                shutil.rmtree(self.output_path.parent)
            os.makedirs(self.output_path.parent)

            kc = self.form_key_get_kc(key)
            self.input_path = Path(dir, "none")
            for _file in file_list:
                self.encrypt_file(kc, _file)
        except Exception as E:
            self.logger.error('\r\n{}'.format(traceback.format_exc()))
            return False
        return True


if __name__ == "__main__":
    kk = Encrypt("E:\\luatools_py3\\tt")
    kk.encrypt_file_list(encode_ch("12345678"), "", ['test.lua'])
    # kc = kk.form_key_get_kc(encode_ch("12345678"))
    # kc_data = b""
    # for i in kc:
    #     print(i.value)
    #     kc_data = kc_data + struct.pack("I", i.value)
